/*
 * Copyright (c) 2024 endless-sky
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "uboot_env_manager.h"
#include "utils_log.h"
#include <cstdlib>
#include <fstream>
#include <sstream>
#include <algorithm>
#include <regex>
#include <array>
#include <memory>

namespace El {
namespace Upgrade {

UBootEnvManager::UBootEnvManager()
{
    // Test if fw_printenv is available
    int result = std::system("which fw_printenv > /dev/null 2>&1");
    if (result != 0) {
        EL_ERROR("fw_printenv command not found, U-Boot environment management disabled");
        fw_available_ = false;
        return;
    }

    fw_available_ = true;
    EL_INFO("U-Boot environment manager initialized with fw_printenv/fw_setenv");
}

UBootEnvManager::~UBootEnvManager() {}

std::string UBootEnvManager::GetEnvVar(const std::string &name)
{
    if (!fw_available_) {
        EL_ERROR("fw_printenv not available");
        return "";
    }

    std::string cmd = "fw_printenv " + name + " 2>/dev/null";
    std::array<char, 128> buffer;
    std::string result;

    FILE *pipe = popen(cmd.c_str(), "r");
    if (!pipe) {
        EL_ERROR("Failed to execute fw_printenv for variable: {}", name);
        return "";
    }

    while (fgets(buffer.data(), buffer.size(), pipe) != nullptr) {
        result += buffer.data();
    }

    int exit_code = pclose(pipe);
    if (exit_code != 0) {
        EL_DEBUG("Variable {} not found", name);
        return "";
    }

    // Parse output: "name=value"
    size_t eq_pos = result.find('=');
    if (eq_pos != std::string::npos) {
        std::string value = result.substr(eq_pos + 1);
        // Remove trailing newline
        if (!value.empty() && value.back() == '\n') {
            value.pop_back();
        }
        EL_DEBUG("Got env var: {}={}", name, value);
        return value;
    }

    return "";
}

bool UBootEnvManager::SetEnvVar(const std::string &name, const std::string &value)
{
    if (!fw_available_) {
        EL_ERROR("fw_setenv not available");
        return false;
    }

    // 对value进行shell转义，避免特殊字符被解释
    std::string escaped_value = "'" + value + "'";
    std::string cmd = "fw_setenv " + name + " " + escaped_value + " 2>/dev/null";
    int result = std::system(cmd.c_str());
    if (result != 0) {
        EL_ERROR("Failed to set environment variable: {}={}", name, value);
        return false;
    }

    EL_DEBUG("Set env var: {}={}", name, value);
    return true;
}

bool UBootEnvManager::DeleteEnvVar(const std::string &name)
{
    if (!fw_available_) {
        EL_ERROR("fw_setenv not available");
        return false;
    }

    std::string cmd = "fw_setenv " + name + " 2>/dev/null";
    int result = std::system(cmd.c_str());

    if (result != 0) {
        EL_ERROR("Failed to delete environment variable: {}", name);
        return false;
    }

    EL_DEBUG("Deleted env var: {}", name);
    return true;
}

std::map<std::string, std::string> UBootEnvManager::GetAllEnvVars()
{
    std::map<std::string, std::string> vars;

    if (!fw_available_) {
        EL_ERROR("fw_printenv not available");
        return vars;
    }

    std::string cmd = "fw_printenv 2>/dev/null";
    std::array<char, 128> buffer;
    std::string result;

    FILE *pipe = popen(cmd.c_str(), "r");
    if (!pipe) {
        EL_ERROR("Failed to execute fw_printenv");
        return vars;
    }

    while (fgets(buffer.data(), buffer.size(), pipe) != nullptr) {
        result += buffer.data();
    }

    pclose(pipe);

    // Parse each line: "name=value"
    std::istringstream iss(result);
    std::string line;
    while (std::getline(iss, line)) {
        if (!line.empty()) {
            size_t eq_pos = line.find('=');
            if (eq_pos != std::string::npos) {
                std::string name = line.substr(0, eq_pos);
                std::string value = line.substr(eq_pos + 1);
                vars[name] = value;
            }
        }
    }

    EL_DEBUG("Got {} environment variables", vars.size());
    return vars;
}

std::string UBootEnvManager::BootSlotToString(BootSlot slot)
{
    switch (slot) {
        case BootSlot::SLOT_A:
            return "a";
        case BootSlot::SLOT_B:
            return "b";
        case BootSlot::UNKNOWN:
        default:
            return "";
    }
}

BootSlot UBootEnvManager::StringToBootSlot(const std::string &str)
{
    if (str == "a") {
        return BootSlot::SLOT_A;
    } else if (str == "b") {
        return BootSlot::SLOT_B;
    } else {
        return BootSlot::UNKNOWN;
    }
}

std::string UBootEnvManager::GetBootSlotInfo()
{
    // Create a JSON-like string with boot slot information
    BootSlot current_slot = GetCurrentBootSlot();
    BootSlot inactive_slot = GetInactiveBootSlot();
    BootSlot target_slot = GetInactiveBootSlot(); // In upgrade context, target is the inactive slot
    BootStatus boot_status = GetBootStatus();
    int boot_count = GetBootCount();
    int max_boot_count = GetMaxBootCount();

    std::string status_str;
    switch (boot_status) {
        case BootStatus::SUCCESSFUL:
            status_str = "SUCCESSFUL";
            break;
        case BootStatus::FAILED:
            status_str = "FAILED";
            break;
        case BootStatus::UNKNOWN:
        default:
            status_str = "UNKNOWN";
            break;
    }

    // Format as a simple JSON string
    std::string json_info = "{\"current_slot\":\"" + BootSlotToString(current_slot) + "\",\"inactive_slot\":\"" +
                            BootSlotToString(inactive_slot) + "\",\"target_slot\":\"" + BootSlotToString(target_slot) +
                            "\",\"boot_status\":\"" + status_str + "\"";

    if (boot_count >= 0) {
        json_info += ",\"boot_count\":" + std::to_string(boot_count);
    }

    if (max_boot_count >= 0) {
        json_info += ",\"max_boot_count\":" + std::to_string(max_boot_count);
    }

    json_info += "}";

    return json_info;
}

BootSlot UBootEnvManager::GetCurrentBootSlot()
{
    std::string slot_str = GetEnvVar(ENV_BOOT_SLOT);
    return StringToBootSlot(slot_str);
}

BootSlot UBootEnvManager::GetInactiveBootSlot()
{
    BootSlot current = GetCurrentBootSlot();
    switch (current) {
        case BootSlot::SLOT_A:
            return BootSlot::SLOT_B;
        case BootSlot::SLOT_B:
            return BootSlot::SLOT_A;
        case BootSlot::UNKNOWN:
        default:
            return BootSlot::SLOT_A; // 默认返回A
    }
}

bool UBootEnvManager::SetBootSlot(BootSlot slot)
{
    if (slot == BootSlot::UNKNOWN) {
        EL_ERROR("Cannot set boot slot to UNKNOWN");
        return false;
    }

    std::string slot_str = BootSlotToString(slot);
    bool success = SetEnvVar(ENV_BOOT_SLOT, slot_str);

    if (success) {
        EL_INFO("Set boot slot to: {}", slot_str);
    }

    return success;
}

bool UBootEnvManager::SwitchBootSlot(BootSlot slot)
{
    if (slot == BootSlot::UNKNOWN) {
        EL_ERROR("Cannot switch boot slot to UNKNOWN");
        return false;
    }

    std::string slot_str = BootSlotToString(slot);
    bool success = SetEnvVar(ENV_BOOT_SLOT, slot_str);

    if (success) {
        EL_INFO("Switched boot slot to: {}", slot_str);
    }

    return success;
}

bool UBootEnvManager::MarkBootSuccessful()
{
    bool success = SetEnvVar(ENV_BOOT_SUCCESSFUL, "1");

    if (success) {
        EL_INFO("Marked boot as successful");
        // 同时重置启动计数器
        ResetBootCount();
    }

    return success;
}

bool UBootEnvManager::ResetBootCount()
{
    bool success = SetEnvVar(ENV_BOOT_COUNT, "0");

    if (success) {
        EL_DEBUG("Reset boot count to 0");
    }

    return success;
}

int UBootEnvManager::GetBootCount()
{
    std::string count_str = GetEnvVar(ENV_BOOT_COUNT);
    if (count_str.empty()) {
        return -1;
    }

    try {
        return std::stoi(count_str);
    } catch (const std::exception &e) {
        EL_ERROR("Invalid boot count value: {}", count_str);
        return -1;
    }
}

int UBootEnvManager::GetMaxBootCount()
{
    std::string count_str = GetEnvVar(ENV_MAX_BOOT_COUNT);
    if (count_str.empty()) {
        return -1;
    }

    try {
        return std::stoi(count_str);
    } catch (const std::exception &e) {
        EL_ERROR("Invalid max boot count value: {}", count_str);
        return -1;
    }
}

bool UBootEnvManager::SetMaxBootCount(int count)
{
    if (count <= 0) {
        EL_ERROR("Max boot count must be positive");
        return false;
    }

    bool success = SetEnvVar(ENV_MAX_BOOT_COUNT, std::to_string(count));

    if (success) {
        EL_INFO("Set max boot count to: {}", count);
    }

    return success;
}

BootStatus UBootEnvManager::GetBootStatus()
{
    std::string status_str = GetEnvVar(ENV_BOOT_SUCCESSFUL);

    if (status_str == "1") {
        return BootStatus::SUCCESSFUL;
    } else if (status_str == "0") {
        return BootStatus::FAILED;
    } else {
        return BootStatus::UNKNOWN;
    }
}

void UBootEnvManager::PrintBootInfo()
{
    BootSlot current_slot = GetCurrentBootSlot();
    BootSlot inactive_slot = GetInactiveBootSlot();
    BootStatus boot_status = GetBootStatus();
    int boot_count = GetBootCount();
    int max_boot_count = GetMaxBootCount();

    EL_INFO("=== Boot Partition Information ===");
    EL_INFO("Current boot slot: {}", BootSlotToString(current_slot));
    EL_INFO("Inactive boot slot: {}", BootSlotToString(inactive_slot));

    std::string status_str;
    switch (boot_status) {
        case BootStatus::SUCCESSFUL:
            status_str = "SUCCESSFUL";
            break;
        case BootStatus::FAILED:
            status_str = "FAILED";
            break;
        case BootStatus::UNKNOWN:
        default:
            status_str = "UNKNOWN";
            break;
    }
    EL_INFO("Boot status: {}", status_str);

    if (boot_count >= 0) {
        EL_INFO("Boot count: {}", boot_count);
    } else {
        EL_INFO("Boot count: NOT_SET");
    }

    if (max_boot_count >= 0) {
        EL_INFO("Max boot count: {}", max_boot_count);
    } else {
        EL_INFO("Max boot count: NOT_SET");
    }

    EL_INFO("==================================");
}

bool UBootEnvManager::ValidateBootConfig()
{
    // 检查必要的环境变量是否存在
    BootSlot current_slot = GetCurrentBootSlot();
    if (current_slot == BootSlot::UNKNOWN) {
        EL_WARN("Boot slot not configured, attempting to set default value");

        // 尝试设置默认的启动槽位为 'a'
        if (SetEnvVar(ENV_BOOT_SLOT, "a")) {
            EL_INFO("Successfully set default boot_slot to 'a'");
            current_slot = BootSlot::SLOT_A;
        } else {
            EL_ERROR("Boot slot not configured properly and cannot be fixed");
            return false;
        }
    }

    int max_boot_count = GetMaxBootCount();
    if (max_boot_count <= 0) {
        EL_WARN("Max boot count not set, using default value 3");
        SetMaxBootCount(3);
    }

    // 检查启动状态
    BootStatus status = GetBootStatus();
    if (status == BootStatus::UNKNOWN) {
        EL_WARN("Boot status unknown, marking as successful");
        MarkBootSuccessful();
    }

    return true;
}

bool UBootEnvManager::ForceBootSlot(BootSlot slot)
{
    if (slot == BootSlot::UNKNOWN) {
        EL_ERROR("Cannot force boot to UNKNOWN slot");
        return false;
    }

    EL_WARN("Force switching to boot slot: {}", BootSlotToString(slot));

    // 设置启动分区
    if (!SetBootSlot(slot)) {
        return false;
    }

    // 重置计数器和状态
    ResetBootCount();
    SetEnvVar(ENV_BOOT_SUCCESSFUL, "0");

    EL_WARN("Emergency boot slot switch completed. System will boot from slot {} on next reboot",
             BootSlotToString(slot));

    return true;
}

} // namespace Upgrade
} // namespace El